from deployment.aliengo_robot import Policy
from deployment.aliengo_robot import AliengoRobot
import time
import numpy as np
import torch
import os

import ctypes
import multiprocessing
from multiprocessing import Process, Queue
from datetime import datetime

import socket
import cv2

def GetAction(shared_obs_base, shared_act_base, control_finish_base,reset_finish_base):
    obs = np.frombuffer(shared_obs_base, dtype=ctypes.c_double)
    act = np.frombuffer(shared_act_base, dtype=ctypes.c_double)
    control_finish = np.frombuffer(control_finish_base, dtype=ctypes.c_double)
    reset_finish = np.frombuffer(reset_finish_base, dtype=ctypes.c_double)
    
    policy = Policy()
    # ================= begin init depth camera ===========================
    mySocket = socket.socket(socket.AF_INET,socket.SOCK_DGRAM)
    mySocket.bind(('localhost',10001))
    mySocket.settimeout(0.01)
    while True:
        try:
            data, _ = mySocket.recvfrom(4800)
            break
        except socket.timeout:
            continue
    depth_image = cv2.imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_GRAYSCALE)
    # =================== end init depth camera ===========================                  
    tt = time.time()
    print('begin get action:')
    #obs prompt\
    for i in range(2):
        policy.GetAction(obs, depth_image)
    # policy.trajectory_history[0] = 0
    print('get first action time:', time.time() - tt)


    while(reset_finish[0] <= 0):
        time.sleep(0.00001)

    last_time = time.time()
    last_last_time = last_time
    step = 0
    states = np.zeros((1000,48))
    actions = np.zeros((1000, 12))
    while True:
        while(last_time + 0.020 - time.time() > 0.):
            time.sleep(0.00001)
        
        
        try:
            data, _ = mySocket.recvfrom(4800)
            depth_image = cv2.imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_GRAYSCALE)
            print('use new image')
            print(last_time - last_last_time)
            last_last_time = last_time
            
        except socket.timeout:
            print('receive image timeout, use saved image')
            pass
        last_time = time.time()
        act[:] = policy.GetAction(obs, depth_image)
        # print('get action time:', time.time()-last_time)
        states[step] = obs
        actions[step] = act
        step += 1
        # print(act)
        if(control_finish[0] > 0.):
            break
    tag = datetime.now().strftime('%b%d_%H-%M-%S')
    path = './data/' + tag
    os.makedirs(path)
    np.save(path + '/state.npy', states[:step])
    np.save(path + '/action.npy', actions[:step])


def RobotControl(shared_obs_base, shared_act_base, control_finish_base,reset_finish_base, control_loop_time,):
    time.sleep(10)
    obs = np.frombuffer(shared_obs_base, dtype=ctypes.c_double)
    act = np.frombuffer(shared_act_base, dtype=ctypes.c_double)
    control_finish = np.frombuffer(control_finish_base, dtype=ctypes.c_double)
    reset_finish = np.frombuffer(reset_finish_base, dtype=ctypes.c_double)
    robot = AliengoRobot()

    #reset the robot
    default_motor_angles = robot.default_dof_pos
    current_motor_angles = robot.GetMotorAngles()
    for t in range(2000):
        blend_ratio = min(t / 1000, 1)
        action = blend_ratio * default_motor_angles + (
                1 - blend_ratio) * current_motor_angles
        # robot._StepInternal(action,process_action=False)
        time.sleep(0.001)
        obs[:] = robot.ConstructObservation()
    
    robot._velocity_estimator.reset()
    robot._state_action_counter = 0
    robot._step_counter = 0
    robot._last_reset_time = time.time()
    robot._init_complete = True

    print('reset finish, begin control')
    reset_finish[0] = 1
    
    #begin control
    begin_control_time = time.time()
    last_time = time.time()
    last_last_time = last_time
    while(True):
        # for _ in range(4):
        # print('step time: ',last_time - last_last_time)
        # last_last_time = last_time
        # last_time = time.time()
        # robot._StepInternal(act)
            # print(act)
            # time.sleep(0.002)
        obs[:] = robot.ConstructObservation()

        control_loop_percent = (time.time() - begin_control_time) / control_loop_time
        # robot.command[0] = 0.4 + 0.0 * control_loop_percent

        if(control_loop_percent > 2.0):
            control_finish[0] = 1
            break


if __name__ == "__main__":
    torch.multiprocessing.set_start_method('spawn')
    shared_act_base = multiprocessing.Array(
        ctypes.c_double, 12, lock=False)
    shared_obs_base = multiprocessing.Array(
        ctypes.c_double, 48, lock=False)
        
    control_finish_base = multiprocessing.Array(
        ctypes.c_double, 1, lock=False)
    reset_finish_base = multiprocessing.Array(
        ctypes.c_double, 1, lock=False)
    
    # depth_image = torch.zeros((60,80))
    # depth_image.share_memory_()


    obs = np.frombuffer(shared_obs_base, dtype=ctypes.c_double)
    act = np.frombuffer(shared_act_base, dtype=ctypes.c_double)
    act[:] = np.zeros(12)


    p1 = Process(target=RobotControl, args=(shared_obs_base, shared_act_base, control_finish_base,reset_finish_base, 10,))
    p2 = Process(target=GetAction, args=(shared_obs_base, shared_act_base, control_finish_base,reset_finish_base, ))

    p1.start()
    p2.start()

    os.system("sudo renice -20 -p %s " % (p1.pid))
    os.system("sudo renice -20 -p %s " % (p2.pid))

    p1.join()
    p2.join()

